#include "multiplexer.h"

MultiplexServer::MultiplexServer(EventLoop *loop,
                                 const InetAddress &listenAddr,
                                 const InetAddress &serverAddr,
                                 int numThreads)
    : _server(loop, listenAddr),
      _backendClient(loop, serverAddr),
      _numTheads(numThreads),
      _oldCounter(0),
      _startTime(Timestamp::now())
{
    _server.setConnectionCallback(std::bind(&MultiplexServer::onClientConnection, this, _1));
    _server.setMessageCallback(std::bind(&MultiplexServer::onClientMessage, this, _1, _2, _3));
    _server.setThreadNum(_numTheads);

    _backendClient.setConnectionCallback(std::bind(&MultiplexServer::onBackendConnection, this, _1));
    _backendClient.setMessageCallback(std::bind(&MultiplexServer::onBackendMessage, this, _1, _2, _3));

    _backendClient.enableRetry();

    // loop->runEvery(10.0, std::bind(&MultiplexServer::printStatistics, this));
}

void MultiplexServer::start()
{
    LOG_INFO << "starting " << _numTheads << " threads.";
    _backendClient.connect();
    _server.start();
}

void MultiplexServer::onClientConnection(const TcpConnectionPtr &conn)
{
    LOG_TRACE << "Client " << conn->peerAddress().toHostPort() << " -> "
              << conn->localhostAddress().toHostPort() << " is "
              << (conn->connected() ? "UP" : "DOWN");
    //取ID
    if (conn->connected())
    {
        int id = -1;
        {
            MutexLockGuard guard(_mutex);
            if (!_availIds.empty())
            {
                id = std::move(_availIds.front());
                _availIds.pop();
                _clientConns[id] = conn;
            }
        }
        //未取到ID
        if (id <= 0)
        {
            conn->shutdown();
        }
        else
        {
            conn->setContext(id);
            char buffer[256];
            snprintf(buffer, sizeof(buffer), "CONN %d FROM %s IS UP\r\n",
                     id,
                     conn->peerAddress().toHostPort().c_str());
            sendBackendString(0, buffer);
        }
    }
    //归还
    else
    {
        if (!conn->getContext().empty())
        {
            int id = boost::any_cast<int>(conn->getContext());
            assert(0 < id && id <= kMaxConns);
            char buffer[256];
            snprintf(buffer, sizeof(buffer), "CONN %d FROM %s IS DOWN\r\n",
                     id, conn->peerAddress().toHostPort().c_str());
            sendBackendString(0, buffer);

            {
                MutexLockGuard guard(_mutex);
                if (_backendConn) //服务器存在
                {
                    _availIds.emplace(id);
                    _clientConns.erase(id);
                }
                else
                {
                    assert(_availIds.empty());
                    assert(_clientConns.empty());
                }
            }
        }
    }
}
void MultiplexServer::onClientMessage(const TcpConnectionPtr &conn, Buffer *buffer, Timestamp)
{
    LOG_DEBUG;
    size_t len_message = buffer->readableBytes();
    //统计
    _transferred.addAndGet(len_message);
    _receivedMessages.incrementAndGet();

    if (!conn->getContext().empty())
    {
        int id = boost::any_cast<int>(conn->getContext());
        sendBackendBuffer(id, buffer);
    }
    else
    {
        buffer->retrieveAll();
    }
    LOG_DEBUG;
}

void MultiplexServer::onBackendConnection(const TcpConnectionPtr &conn)
{
    LOG_TRACE << "Backend " << conn->localhostAddress().toHostPort() << " -> "
              << conn->peerAddress().toHostPort() << " is "
              << (conn->connected() ? "UP" : "DOWN");
    std::vector<TcpConnectionPtr> connsToDestroy;
    if (conn->connected())
    {
        {
            MutexLockGuard guard(_mutex);
            _backendConn = conn;
            //初始化ID队列
            assert(_availIds.empty());
            for (int i = 1; i <= kMaxConns; ++i)
            {
                _availIds.push(i);
            }
        }
    }
    else
    {
        //释放conn
        {
            MutexLockGuard guard(_mutex);
            _backendConn.reset();
            connsToDestroy.reserve(_clientConns.size());
            for (auto it = _clientConns.begin(); it != _clientConns.end(); ++it)
            {
                connsToDestroy.emplace_back(it->second);
            }
            //释放
            _clientConns.clear();
            while (!_availIds.empty())
            {
                _availIds.pop();
            }
        }
    }
    for (auto &item : connsToDestroy)
    {
        item->shutdown();
    }
}
void MultiplexServer::onBackendMessage(const TcpConnectionPtr &conn, Buffer *buffer, Timestamp)
{
    LOG_DEBUG;
    size_t len_message = buffer->readableBytes();
    _transferred.addAndGet(len_message);
    _receivedMessages.incrementAndGet();
    sendToClient(buffer);
    LOG_DEBUG;
}
void MultiplexServer::sendToClient(Buffer *buffer)
{
    LOG_DEBUG;
    //取出header
    while (buffer->readableBytes() > kHeaderLen)//至少得知道ID
    {
        int len_message = static_cast<uint8_t>(*buffer->peek());
        if (buffer->readableBytes() < kHeaderLen + len_message)
        {
            break;
        }
        else
        {
            int id = static_cast<uint8_t>(buffer->peek()[1]);
            id |= (static_cast<uint8_t>(buffer->peek()[2]) << 8);
            if (id != 0)
            {
                LOG_DEBUG<<id;
                TcpConnectionPtr clientConn;
                {
                    MutexLockGuard guard(_mutex);//多线程退化为单线程
                    auto it = _clientConns.find(id);
                    if (it != _clientConns.end())
                    {
                        clientConn = it->second;
                    }
                }
                if (clientConn)
                {
                    LOG_DEBUG<<buffer->peek();
                    clientConn->send(buffer->peek() + kHeaderLen, len_message);
                }
            }
            else
            {
                string cmd(buffer->peek() + kHeaderLen, len_message);
                LOG_INFO << "Backend cmd " << cmd;
                doCommand(cmd);
            }
        }
        buffer->retrieve(kHeaderLen + len_message);
    }
    LOG_DEBUG;
}
void MultiplexServer::doCommand(const string &cmd)
{
    TcpConnectionPtr clientConn;
    static const string kDisconnectCmd = "DISCONNECT ";

    if (cmd.size() > kDisconnectCmd.size() && std::equal(kDisconnectCmd.begin(), kDisconnectCmd.end(), cmd.begin()))
    {
        int connId = atoi(&cmd[kDisconnectCmd.size()]);

        {
            MutexLockGuard guard(_mutex);
            std::map<int, TcpConnectionPtr>::iterator it = _clientConns.find(connId);
            if (it != _clientConns.end())
            {
                clientConn = it->second;
            }
        }
    }
    clientConn->shutdown();
}

void MultiplexServer::printStatistics()
{
    Timestamp endTime = Timestamp::now();
    int64_t newCounter = _transferred.get();
    int64_t bytes = newCounter - _oldCounter;
    int64_t msgs = _receivedMessages.getAndSet(0);
    double time = timeDifference(endTime, _startTime);
    printf("%4.3f MiB/s %4.3f Ki Msgs/s %6.2f bytes per msg\n",
           static_cast<double>(bytes) / time / 1024 / 1024,
           static_cast<double>(msgs) / time / 1024,
           static_cast<double>(bytes) / static_cast<double>(msgs));

    _oldCounter = newCounter;
    _startTime = endTime;
}

/*
 * 封包及发送
 */
//加包首
void MultiplexServer::sendBackendPacket(int id, Buffer *buffer)
{
    size_t len_message = buffer->readableBytes();
    assert(len_message <= kMaxPacketLen); // len_message 小于8位可以表达的数字,0-255.

    uint8_t header[kHeaderLen] = // 3个字节
        {
            static_cast<uint8_t>(len_message),
            // ID可以0-65535，2字节，4个十六进制。
            static_cast<uint8_t>(id & 0xFF),
            static_cast<uint8_t>((id & 0xFF00) >> 8)};

    buffer->prepend(header, kHeaderLen);

    TcpConnectionPtr backendConn; // copy on write
    {
        MutexLockGuard guard(_mutex);
        backendConn = _backendConn;
    }
    if (backendConn)
    {
        backendConn->send(buffer);
    }
}
void MultiplexServer::sendBackendString(int id, const string &msg)
{
    assert(msg.size() <= kMaxPacketLen);
    Buffer buffer;
    buffer.append(msg);
    sendBackendPacket(id, &buffer);
}
void MultiplexServer::sendBackendBuffer(int id, Buffer *buffer)
{
    while (buffer->readableBytes() > kMaxPacketLen)
    {
        Buffer packet;
        packet.append(buffer->peek(), kMaxPacketLen);
        buffer->retrieve(kMaxPacketLen);
        sendBackendPacket(id, &packet);
    }
    if (buffer->readableBytes() > 0)
    {
        sendBackendPacket(id, buffer);
    }
}
